import subprocess
import rrd_base
import re
import time

import snmp
import color

try:
    from pyrrd.rrd import DataSource, RRA, RRD
    from pyrrd.graph import DEF, CDEF, VDEF
    from pyrrd.graph import LINE, AREA, GPRINT
    from pyrrd.graph import ColorAttributes, Graph
except:
    print "Couldn't load required modules. You probably need to install it with:"
    sys.exit('sudo apt-get install python-pyrrd python-pysnmp4')

class RrdMemory(rrd_base.RrdBase):
    def __init__(self, domain):
        super(RrdMemory,self).__init__(domain=domain, group="cpu", instance="memory")

    def appendDataSources(self, dataSources):
        assert 0, "RrdMemory: appendDataSources() should not be called"

    def create_file(self):
        dataSources = []
        self.appendDataSources(dataSources)

        roundRobinArchives = []
        for func in ['AVERAGE', 'MAX', 'MIN']:
            for step,row in self.rra_pattern():
                roundRobinArchives.append(RRA(cf=func, xff=0.5, steps=step, rows=row))

        self.myRRD = RRD(self.filename, ds=dataSources, rra=roundRobinArchives, step=self.data_step())
        self.myRRD.create()

class RrdMemoryLocal(RrdMemory):
    def __init__(self, domain):
        super(RrdMemoryLocal,self).__init__(domain=domain)

    def appendDataSources(self, dataSources):
        for ds in ("ram_used", "ram_free", "ram_shared", "ram_buffers", "ram_cached", "swap_used", "swap_free"):
            dataSources.append(DataSource( dsName=ds,    dsType='GAUGE', heartbeat=self.heartbeat()))

    def update(self):
        sp = subprocess.Popen(["/usr/bin/free", "-b"], shell=False, stdout=subprocess.PIPE)
        lines = sp.communicate()[0].split("\n")
        mem = lines[1].split()
        swap= lines[3].split()

        self.myRRD.bufferValue(int(round(time.time())),
                               int(mem[2]),
                               int(mem[3]),
                               int(mem[4]),
                               int(mem[5]),
                               int(mem[6]),
                               int(swap[2]),
                               int(swap[3]))
        self.myRRD.update()

    def graph(self):
        def_ram_used     = DEF(rrdfile=self.myRRD.filename, vname='ram_used',    cdef="AVERAGE", dsName='ram_used')
        def_ram_free     = DEF(rrdfile=self.myRRD.filename, vname='ram_free',    cdef="AVERAGE", dsName='ram_free')
        def_ram_shared   = DEF(rrdfile=self.myRRD.filename, vname='ram_shared',  cdef="AVERAGE", dsName='ram_shared')
        def_ram_buffers  = DEF(rrdfile=self.myRRD.filename, vname='ram_buffers', cdef="AVERAGE", dsName='ram_buffers')
        def_ram_cached   = DEF(rrdfile=self.myRRD.filename, vname='ram_cached',  cdef="AVERAGE", dsName='ram_cached')
        def_swap_used    = DEF(rrdfile=self.myRRD.filename, vname='swap_used',  cdef="AVERAGE", dsName='swap_used')
        def_swap_free    = DEF(rrdfile=self.myRRD.filename, vname='swap_free',  cdef="AVERAGE", dsName='swap_free')
        cdef_ram_used    = CDEF(vname='ram_in_use', rpn='{0},{1},-,{2},-'.format(def_ram_used.vname, def_ram_buffers.vname, def_ram_cached.vname))
        cdef_total_ram   = CDEF(vname="total_ram",  rpn='{0},{1},+'.format(def_ram_used.vname, def_ram_free.vname))
        area_ram_used    = AREA(defObj=cdef_ram_used,     color=color.RGBA['red'],         legend='RAM Used',    stack=True)
        area_ram_free    = AREA(defObj=def_ram_free,      color=color.RGBA['forestgreen'], legend='RAM Free',    stack=True)
        area_ram_buffers = AREA(defObj=def_ram_buffers,   color=color.RGBA['green'],       legend='RAM Buffers', stack=True)
        area_ram_cached  = AREA(defObj=def_ram_cached,    color=color.RGBA['darkgreen'],   legend='RAM Cached',  stack=True)
        area_ram_shared  = AREA(defObj=def_ram_shared,    color=color.RGBA['purple'],      legend='RAM Shared',  stack=True)
        area_swap_used   = AREA(defObj=def_swap_used,     color=color.RGBA['red'],         legend='Swap Used',   stack=True)
        area_swap_free   = AREA(defObj=def_swap_free,     color=color.RGBA['green'],       legend='Swap Free',   stack=True)
        line_total_ram   = LINE(3, defObj=cdef_total_ram, color=color.RGBA['black'])

        for tf in self.graph_time_frames().keys():
            g = Graph(self.image_name(tf),
                      start=int(time.time())-self.graph_time_frames()[tf], end=int(time.time()),
                      lower_limit=0.000001,
                      base=1024,
                      title='"Memory Usage for {0} - Updated {1}"'.format(self.data_domain, time.strftime("%Y%m%d at %H:%M")),
                      vertical_label='Bytes',
                      height=self.graph_height, width=self.graph_width,
                      x_grid=self.grid[tf],
                      step=self.graph_step_times[tf])
            g.data.extend([def_ram_used, def_ram_free, def_ram_shared, def_ram_buffers, def_ram_cached, def_swap_used, def_swap_free, cdef_ram_used, cdef_total_ram, area_ram_free, area_ram_buffers, area_ram_cached, area_ram_shared, area_ram_used, area_swap_used, area_swap_free, line_total_ram])
            g.write()

            if False:
                g = Graph(self.icon_name(tf),
                          start=int(time.time())-self.graph_time_frames()[tf], end=int(time.time()),
                          lower_limit=0.000001,
                          base=1024,
                          no_legend=True,
                          height=self.icon_height, width=self.icon_width,
                          step=self.graph_step_times[tf])
                g.data.extend([def_ram_used, def_ram_free, def_ram_shared, def_ram_buffers, def_ram_cached, def_swap_used, def_swap_free, cdef_ram_used, cdef_total_ram, area_ram_free, area_ram_buffers, area_ram_cached, area_ram_shared, area_ram_used, area_swap_used, area_swap_free, line_total_ram])
                g.write()

class RrdMemorySnmp(RrdMemory):
    def __init__(self, domain):
        super(RrdMemorySnmp,self).__init__(domain=domain)

    def appendDataSources(self, dataSources):
        for ds in ["ram_total", "ram_used"]:
            dataSources.append(DataSource( dsName=ds, dsType='GAUGE', heartbeat=self.heartbeat()))

    def update(self):
        used = snmp.get((1,3,6,1,2,1,25,2,3,1,6,101))
        total = snmp.get((1,3,6,1,2,1,25,2,3,1,5,101))
        self.myRRD.bufferValue(int(round(time.time())), int(total)*1024, int(used)*1024)
        self.myRRD.update()

    def graph(self):
        d_t     = DEF(rrdfile=self.myRRD.filename, vname='ram_total', reduce="AVERAGE", dsName='ram_total')
        d_u     = DEF(rrdfile=self.myRRD.filename, vname='ram_used',  reduce="AVERAGE", dsName='ram_used')
        c_free  =CDEF(vname='free', rpn="{0},{1},-".format(d_t.vname, d_u.vname))
        a_u     = AREA(defObj=d_u,    color=color.RGBA['red'], legend='RAM Used',    stack=True)
        a_f     = AREA(defObj=c_free, color=color.RGBA['forestgreen'], legend='RAM Free',    stack=True)

        for tf in self.graph_time_frames().keys():
            g = Graph("{0}/{1}.png".format(self.graph_dir, tf),
                      start=int(time.time())-self.graph_time_frames()[tf],
                      end=int(time.time()),
                      title='"Memory Usage for {0} - Updated {1}"'.format(self.data_domain, time.strftime("%Y%m%d at %H:%M")),
                      vertical_label='Bytes',
                      lower_limit=1,
                      height=self.graph_height,
                      width=self.graph_width,
                      base=1024,
                      x_grid=self.grid[tf],
                      step=self.graph_step_times[tf])

            g.data.extend([d_t, d_u, c_free, a_u, a_f])

            g.write()
